# coding=utf-8
import os


"""
注意事项：
    1. 只检测 fq 后缀的文件名称
    2. 名称对应关系表名称为 name_pair.csv，第一列是旧名字，第二列是新名字
    3. 新名称的文件已存在会被覆盖
    4. 如果有异常，会生成一个 error_log.txt 文件
    5. 数据必须跟此工具一个目录

    3. fq要求：
        a. 由 【Field 1】-【Product Code】-【Sample Type】-【Field 2】_【R1|R2】.【fastq.gz|fq.gz】 组成
        b. Field 1:  建议使用患者姓名缩写
        c. Product Code: 产品代号，如 ONCO13
        d. Sample Type: 样本类型，DNA, RNA, tumor, or normal
        e. Field 4: 通常是医院内使用的病理编号或公司定制的样本编号
        f. 以上4个字段都只允许似乎用字母和数字命名

        fq 名称完整示例： ZS-ONCO13-DNA-ABC20230101_R1.fastq.gz

"""

fq_require = '''
fq要求：
    a. 由 【Field 1】-【Product Code】-【Sample Type】-【Field 2】_【R1|R2】.【fastq.gz|fq.gz】 组成
    b. Field 1:  建议使用患者姓名缩写
    c. Product Code: 产品代号，如 ONCO13
    d. Sample Type: 样本类型，DNA, RNA, tumor, or normal
    e. Field 4: 通常是医院内使用的病理编号或公司定制的样本编号
    f. 以上4个字段都只允许似乎用字母和数字命名

    fq 名称完整示例： ZS-ONCO13-DNA-ABC20230101_R1.fastq.gz
'''

def filename_err(f_name):
    """ 检查文件格式，不合要求的返回 False """
    if f_name.count('-') != 3:
        return True
    f_name = f_name.replace('.fastq.gz', '').replace('.fq.gz', '')
    if '_R1' not in f_name and '_R2' not in f_name:
        return True
    if f_name.count('_') != 1:
        return True
    f_name = f_name.replace('_R1', '').replace('_R2', '')
    for w in f_name.split('-'):
        if not w.isalnum():
            return True


name_pair = 'name_pair.csv'
error_log = 'error_log.txt'

# 文件存在的话，就直接使用其进行数据改名和合并
err_info = ''
if os.path.isfile(name_pair):
    # 获取名称对应关系
    names = {}
    with open(name_pair) as fi:
        for line in fi:
            if not line.strip():
                continue
            cells = line.strip().split(',')
            if not cells[1]:
                err_info = f'{err_info}文件没有新名字 {cells[0]}\n'
                continue
            if filename_err(cells[1]):
                err_info = f'{err_info}名字不规范 {cells[1]}\n'
                continue
            if not cells[1].endswith('.fq.gz') and not cells[1].endswith('.fastq.gz'):
                cells[1] = f'{cells[1]}.fq.gz'
            if cells[1] in names:
                names[cells[1]].add(cells[0])
            else:
                names[cells[1]] = set([cells[0]])
            if os.path.isfile(cells[1]):
                err_info = f'{err_info}文件已存在 {cells[1]}\n'
                continue
    for new_name, fqs in names.items():
        fqs = sorted(list(fqs))  #! 不排序合并的 R1 和 R2 可能不配对了
        for fq in fqs:
            if not os.path.isfile(fq):
                err_info = f'{err_info}文件不存在 {fq}\n'
        # 文件找不到就下一个了
        if err_info:
            continue
        if len(fqs) == 1:
            os.rename(fqs[0], new_name)
            print(f'{fqs[0]} 重命名为  {new_name}')
        else:
            file_content = b''
            for fq in fqs:
                with open(fq, 'rb') as fi:
                    file_content += fi.read()
            with open(new_name, 'wb') as fo:
                fo.write(file_content)
            print(f'{fqs} 合并为  {new_name}')
# 文件不存在就创建一个，并填入旧文件名
else:
    file_text = ''
    for fq in os.listdir():
        if not fq.endswith('.fq.gz') and not fq.endswith('.fastq.gz'):
            continue
        file_text = f'{file_text}{fq},\n'
    with open(name_pair, 'w', encoding='utf-8') as fo:
        fo.write(file_text)



if err_info:
    if '名字不规范' in err_info:
        err_info += fq_require
    with open(error_log, 'w') as fo:
        fo.write(err_info)
else:
    if os.path.isfile(error_log):
        os.remove(error_log)

